#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation

# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA
#

from numpy import log2, exp, cos, sin, ones, c_
from numpy import linalg, asarray, empty, meshgrid

from scipy import linspace

from utils.Quadrature import LegendreGaussQuadrature
from utils.FemBases import LagrangeP1, LagrangeP2
from utils.ProjOperators import ProjOnShapeFunctions
from utils.FemToReal import FemToReal
from utils.Grid import Grid1D
from utils.Spline import Spline

from models.CDRSolver import CDRSolver
from models.SecondOrderStatic1D import SecondOrderStatic1D
from models.Poisson2DFiniteDifference import Poisson2DFiniteDifference
from models.AssembleStiffnesMassMatrices import AssembleStiffnesMassMatrices

from solvers.DirichletTimeStepper import DirichletTimeStepper

from input_output.DerivativePlots import DerivativePlots
import matplotlib.pyplot as plt
from scipy.sparse.construct import spdiags

import logging
from utils.MainLog import runGlobalLogger

import unittest

(d, c, r) = (-1., 2.0, 0.50)


class PdeTestSuite(unittest.TestCase):
    """ test elliptic and parabolic solver. """

    @classmethod
    def cinf_solution(cls, points, time=0):  # Test 1
        """ exact solution. """
        return exp(points)

    @classmethod
    def cinf_load(cls, points, time=0):
        """ rhs. """
        return (d + c + r) * exp(points)

    @classmethod
    def exact_solution(cls, points, time=0):  # Test 2
        return points * (3 + points)

    @classmethod
    def gradient(cls, points, time=0):
        return 2 * d * points + d * 3

    @classmethod
    def load(cls, points, time=0):
        return 2 * d * ones(points.shape) + c * (3 + 2 * points) + \
                                        r * points * (3 + points)

    def setUp(self):
        """ I don't know how to avoid this """

        self.m_nsegs = 100
        self.m_xl = -1
        self.m_xr = 1
        self.m_quad = LegendreGaussQuadrature(10)
        self.m_lagrange_p2 = LagrangeP2()
        self.m_lagrange_p1 = LagrangeP1()

        self.m_error_tol = 1.0e-03

        self.m_shape_funct = self.m_lagrange_p1
        self.m_shape_funct = self.m_lagrange_p2

    @staticmethod
    def _comp_errors(mesh, sfuncts, numeric_sol, exact, time=0,
            plot_results = False):
        """ reconstruct to plot """

        recons = FemToReal(mesh, sfuncts)
        solu = recons.reconstruct(numeric_sol)
        y1 = exact(solu[:, 0], time)
        if plot_results:
            from time import sleep
            axis = plt.figure(figsize=(11.5, 5.)).gca()
            axis.plot(solu[:, 0], solu[:, 1], '-')
            axis.plot(solu[:, 0], y1, 'x')
            plt.show()
        return linalg.norm(solu[:, 1].flatten() - y1.flatten())


class Projection(PdeTestSuite):

    """ Test projection operator """

    def runTest(self):
        mesh = Grid1D(self.m_nsegs, self.m_xl, self.m_xr)

        sfuncts = self.m_shape_funct
        numeric_sol = ProjOnShapeFunctions(sfuncts, self.m_quad, mesh,
                              self.exact_solution)
        error_norm = self._comp_errors(mesh, sfuncts, numeric_sol,
                self.exact_solution)
        self.assert_(self.m_error_tol > abs(error_norm))


class AssembledMatrices(PdeTestSuite):

    """   a*u_xx + b*u_x + c*u  = load
          BC : Dirichlet
    """

    def runTest(self):
        """ assemble global matrix, project exact solution, project
            exact solution and compare agains projected rhs.
        """
        pde_solver = SecondOrderStatic1D(self.m_nsegs, self.m_xl, self.m_xr)

        sfuncts = self.m_shape_funct
        pde_solver.set_sfunc_and_quads(sfuncts, self.m_quad)
        (d_is_real, c_is_real, r_is_real) = (True, True, True)
        pde_solver.set_cos(d, c, r, d_is_real, c_is_real, r_is_real)
        pde_solver.set_load(self.load)

        pde_solver.set_bd_conditions(self.exact_solution(self.m_xl),
                                     self.exact_solution(self.m_xr))

        # Get the matrices in diagonal form
        stiff_diag = pde_solver.get_stiff_dformat()
        data = [list(stiff_diag[:, 0].flatten()), list(stiff_diag[:,
                1].flatten()), list(stiff_diag[:, 2].flatten())]
        diags = [0, 1, -1]

        if 3 == sfuncts.dof():
            data.append(list(stiff_diag[:, 3].flatten()))
            data.append(list(stiff_diag[:, 4].flatten()))
            diags.append(2)
            diags.append(-2)

        mat_size = stiff_diag[:, 0].size

        sys = spdiags(asarray(data), asarray(diags), mat_size, mat_size).tocsr()
        rhs = pde_solver.get_load_terms()

        mesh = pde_solver.get_mesh()
        exact_projected = ProjOnShapeFunctions(sfuncts, self.m_quad, mesh,
                self.exact_solution)
        rhs[0] += self.gradient(self.m_xl)
        rhs[-1] -= self.gradient(self.m_xr)

        error_norm = linalg.norm(sys * exact_projected - rhs)
        self.assert_(self.m_error_tol > abs(error_norm))


class SolvePoisson1D(PdeTestSuite):

    """   a*u_xx + b*u_x + c*u  = self.m_load
          BC : Dirichlet
    """

    def runTest(self):
        """ Make sure the static engine is working fine. """

        pde_solver = SecondOrderStatic1D(self.m_nsegs, self.m_xl, self.m_xr)

        sfuncts = self.m_shape_funct
        pde_solver.set_sfunc_and_quads(sfuncts, self.m_quad)

        (d_is_real, c_is_real, r_is_real) = (True, True, True)
        pde_solver.set_cos(d, c, r, d_is_real, c_is_real, r_is_real)
        pde_solver.set_load(self.load)

        pde_solver.set_bd_conditions(self.exact_solution(self.m_xl),
                self.exact_solution(self.m_xr))

        numeric_sol = pde_solver.solve()
        mesh = pde_solver.get_mesh()
        error_norm = self._comp_errors(mesh, sfuncts, numeric_sol,
                self.exact_solution, 0)
        self.assert_(self.m_error_tol > abs(error_norm))


class RateOfConvergenceStatic(PdeTestSuite):
    """ Convergence for lagrange P1 and lagrange P2 methods."""

    def runTest(self):
        print ' '
        conv_log = logging.getLogger('RateOfConvergence')

        nsegs = asarray([int(2 ** k) for k in range(2, 12)], 'int')
        conv_log.info('             exp(x) ')
        conv_log.info('    n          error      log2(ratio)')
        conv_log.info('Lagrange P1')
        (d_is_real, c_is_real, r_is_real) = (True, True, True)
        for sfuncts in [self.m_lagrange_p1, self.m_lagrange_p2]:
            prev = .050
            for k in nsegs:
                pde_solver = SecondOrderStatic1D(k, self.m_xl, self.m_xr)

                pde_solver.set_sfunc_and_quads(sfuncts, self.m_quad)

                pde_solver.set_cos(d, c, r, d_is_real, c_is_real, r_is_real)
                pde_solver.set_load(self.cinf_load)

                pde_solver.set_bd_conditions(self.cinf_solution(self.m_xl),
                        self.cinf_solution(self.m_xr))

                numeric_sol = pde_solver.solve()
                mesh = pde_solver.get_mesh()
                l2_error = self._comp_errors(mesh, sfuncts, numeric_sol,
                        self.cinf_solution, 0)
                conv_log.info(' %5d    %12.11f    %12.11f' % (k,
                                l2_error, log2(prev / l2_error)))
                prev = l2_error
            conv_log.info('')
            conv_log.info('Lagrange P2')


class Diffusion1D(PdeTestSuite):
    """ The most important engine in the GBM pricing suite. """

    def exact_sol(self, x, time=0):
        return exp( (self.m_d + self.m_r) * time) * exp(-x)

    def runTest(self):
        print ' '
        ntime_steps = 400
        time2simulate = .50

        self.m_d = 0
        self.m_r = 1
        self.m_nsegs = 20
        pde_solver = CDRSolver(self.m_nsegs, self.m_xl, self.m_xr)

        pde_solver.set_sfunc_and_quads(self.m_shape_funct, self.m_quad)

        (d_is_real, c_is_real, r_is_real) = (True, True, True)
        pde_solver.set_diffus_cos(self.m_d, 0.0, self.m_r, d_is_real, c_is_real,
                                  r_is_real)

      # pde_solver.set_load(f)

        pde_solver.set_ivalf(self.exact_sol)
        pde_solver.set_lbc(self.exact_sol)
        pde_solver.set_rbc(self.exact_sol)

        time_stepper = DirichletTimeStepper()

        # Use 5000 time steps to get accurancy
        # time_stepper.set_time_step_method('ei')
        time_stepper.set_time_step_method('cn')

        time_stepper.set_dolog(True)
        time_stepper.set_number_steps(ntime_steps)
        time_stepper.set_time2simulate(time2simulate)

        pde_solver.set_time_stepper(time_stepper)

        numeric_sol = pde_solver.solve()
        mesh = pde_solver.get_mesh()

        error_norm = self._comp_errors(mesh, self.m_shape_funct, numeric_sol,
                self.exact_sol, time2simulate, True)
        print error_norm
        self.assert_(1.0e-02 > abs(error_norm))


class Poisson2D(PdeTestSuite):
    """ Test static solver. """

    # 3 test cases, 2 Exact solutions and 1 very smooth

    def exact_solution_2d(self, points):
        """ known solutions. """

        x = points[:, 0].reshape((points[:, 0].size, 1))
        y = points[:, 1].reshape((points[:, 0].size, 1))

        return cos(x)-sin(y)
        #return 0.2 + 0.1 * exp( -0.3 * x)
        #return x** 3
        #return ones(x)

    def load_2d(self, points):
        """ Force term in mechanical engineer. """

        x, y = points[0], points[1]

        return cos(x) - sin(y)
        #return -0.09 * exp( -0.3 * x)
        #return -6 * x
        #return 0

    def _init_solver(self):
        """ gloabal parametes for all test cases. """

        delta_x, delta_y = 10. , 0.05
        delta_x, delta_y = .025 , 0.025
        self.m_p2d = Poisson2DFiniteDifference()

        self.m_p2d.set_simu_domain(-1.5, -1., 1., 1.)
        #self.m_p2d.set_simu_domain(40, 0, 120, 2)
        self.m_p2d.set_xlxr(self.exact_solution_2d,
                                  self.exact_solution_2d)
        self.m_p2d.set_bott_top(self.exact_solution_2d, self.exact_solution_2d)
        self.m_p2d.set_dx_dy(delta_x, delta_y)

    def test_analytical_load(self):
        """ Let the solver know about the load as an analytical expression. """

        self.m_p2d.set_continuos_load(self.load_2d)
        return self.m_p2d.solve()

    def test_discrete_load(self):
        """ Let the solver know about the load at the mesh points. """

        from itertools import izip, count

        (x_mesh, y_mesh) = self.m_p2d.get_mesh()

        # Analytical Load

        (x_coor, Y) = meshgrid(x_mesh, y_mesh)
        discrete_load = empty((x_coor.size, 1))
        for (i, val) in izip(count(), c_[x_coor.flatten(), Y.flatten()]):
            discrete_load[i] = self.load_2d(val)
        discrete_load = discrete_load.reshape((y_mesh.size, x_mesh.size)).T

        self.m_p2d.set_discrete_load(discrete_load)
        return self.m_p2d.solve()

    def runTest(self):
        """
            Solve a simple laplace problem using continuos load and
            discretize load.
        """
        # Analytical Load

        self._init_solver()
        (x_mesh, y_mesh, solu) = self.test_analytical_load()

        # Exact solution

        (x_coor, Y) = meshgrid(x_mesh, y_mesh[::-1])
        analytical = self.exact_solution_2d(c_[x_coor.flatten(), Y.flatten()])

        plot_results = False
        if plot_results:
            dplots = DerivativePlots()
            dplots.set_vol_surface(x_coor, Y, solu,
                                 solu-analytical.reshape(solu.shape))

            dplots.plot_vol_surfaces()

        # a) Numeric vs Exact.

        l2_error = linalg.norm(analytical.flatten() - solu.flatten())
        self.assert_(self.m_error_tol > l2_error)

        # Discrete Load

        self._init_solver()
        (x_mesh, y_mesh, solu) = self.test_discrete_load()

        # b) Numeric vs Exact.

        l2_error = linalg.norm(analytical.flatten() - solu.flatten())
        self.assert_(self.m_error_tol > l2_error)

class Splines(PdeTestSuite):
    """ Variable coeffienents pde solver. """

    def runTest(self):
        x = asarray(linspace(0, self.m_nsegs, self.m_nsegs, False))
        f = self.gradient
        y = f(x)

        my_spline = Spline()
        my_spline.spline(x, y)

        # chk Results

        shift = 0.5
        yt = f(x + shift * ones(x.shape))
        sp_vals = asarray([my_spline.value([x[s] + shift], s) for s in
                       range(x.size - 1)]).flatten()
        error_norm = linalg.norm(yt[:-1] - sp_vals)
        assert self.m_error_tol > error_norm


class ExactVsConstant(PdeTestSuite):
    """ compare the 2 solvers when the coeffients are constant. """

    def runTest(self):
        """
            Compare the pde solver using constant coeffients and spline
            variable coeffients.
        """


        sfuncts = self.m_shape_funct
        mesh = Grid1D(self.m_nsegs, self.m_xl, self.m_xr)

        # constant term PDE

        is_real = True
        ass_const = AssembleStiffnesMassMatrices(mesh)
        ass_const.set_sfunc_and_quads(sfuncts, self.m_quad)
        ass_const.set_cos(d, c, r, is_real, is_real, is_real)
        ass_const.discretize()
        stiff_const = ass_const.get_stiffness()

        # non-constant term PDE

        ass_const = AssembleStiffnesMassMatrices(mesh)
        x = mesh.get_grid()
        dspline = Spline(x, d * ones(x.shape))
        cspline = Spline(x, c * ones(x.shape))
        rspline = Spline(x, r * ones(x.shape))

        is_real = False
        ass_spline = AssembleStiffnesMassMatrices(mesh)
        ass_spline.set_sfunc_and_quads(sfuncts, self.m_quad)
        ass_spline.set_cos(dspline, cspline, rspline, \
                is_real, is_real, is_real)
        ass_spline.discretize()
        stiff_spline = ass_spline.get_stiffness()

        assert stiff_const.shape == stiff_spline.shape
        assert self.m_error_tol > linalg.norm(stiff_const - stiff_spline)


def my_suite():
    """
        Any running test should be added here.
    """
    suite = unittest.TestSuite()

    suite.addTest(Poisson2D())
    suite.addTest(Projection())
    suite.addTest(AssembledMatrices())
    suite.addTest(SolvePoisson1D())
    suite.addTest(Diffusion1D())
    suite.addTest(RateOfConvergenceStatic())

    suite.addTest(Splines())
    suite.addTest(ExactVsConstant())

    return suite


if __name__ == '__main__':

    runGlobalLogger()
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(my_suite())

